from sbi.utils import MultipleIndependent
from sbi.inference import MNLE
from torch.distributions import Uniform, Beta, Binomial, Bernoulli, Categorical, Gamma
import torch
import numpy as np
from typing import Callable

__all__ = [
    "torch_prior",
    "train_mnle",
    "extract_theta_info",
    "gpu_available"
]

def torch_prior(
    thetas_low_boundary: list,
    thetas_high_boundary: list,
    torch_device=None
) -> MultipleIndependent:

    torch_device = gpu_available(torch_device)

    thetas_low_boundary_float = [float(low) for low in thetas_low_boundary]
    thetas_high_boundary_float = [float(high) for high in thetas_high_boundary]

    uniform_distributions = [
        Uniform(
            low=torch.tensor([low], dtype=torch.float32).to(torch_device),
            high=torch.tensor([high], dtype=torch.float32).to(torch_device)
        )
        for low, high in zip(thetas_low_boundary_float, thetas_high_boundary_float)
    ]
    bernoulli_distribution = Bernoulli(
        probs=torch.tensor([0.5], dtype=torch.float32).to(torch_device)
    )
    # 将伯努利分布添加到均匀分布列表中
    all_distributions = uniform_distributions + [bernoulli_distribution]
    # 创建提议分布并将其移动到 GPU
    proposal = MultipleIndependent(all_distributions, validate_args=False)

    return proposal


def gpu_available(torch_device):
    if torch_device is None:
        if torch.cuda.is_available():
            torch_device = torch.device("cuda")
        else:
            torch_device = torch.device("cpu")
    return torch_device

def train_mnle(
    proposal: MultipleIndependent,
    train_data_torch,
    save_name=None,
    torch_device=None
):

    if save_name is not None:
        import pickle
        from pathlib import Path
        save_name = Path(save_name).with_suffix(".pkl")

        if save_name.exists():
            mnle = pickle.load(open(save_name, 'rb'))
            return mnle

    torch_device = gpu_available(torch_device)

    # Initialise the MNLE trainer
    trainer = MNLE(prior=proposal, device=torch_device.type)
    # Add presimulated training data to the MNLE trainer
    trainer = trainer.append_simulations(
        train_data_torch[:, :-2], train_data_torch[:, -2:]
    )
    # Train the network and return the resulting MNLE
    mnle = trainer.train(
        training_batch_size=128 if torch_device.type == "cpu" else 1000
    )

    if save_name is not None:
        with open(save_name, 'wb') as f:
            pickle.dump(mnle, f)

    return mnle


def extract_theta_info(theta_info):
    theta_name = list(theta_info.keys())
    theta_bs = np.array(list(theta_info.values()))
    return theta_name, [theta_bs[:, 0], theta_bs[:, 1]]
